from llama_cpp import Llama
from langchain_community.llms import LlamaCpp

# 初始化模型（改用Q4_K_M量化更稳定）
llm = Llama(
    model_path=r"D:\ideaSpace\MyPython\models\Phi-3-mini-4k-instruct.Q4_K_M.gguf",
    n_threads=4,
    n_ctx=512,
    verbose=False, # 不显示详细日志
    n_batch=512,  # 批处理大小
    n_gpu_layers=0,  # CPU运行（如果有GPU可以设置层数）
)
llm2 = LlamaCpp(
    model_path=r"D:\ideaSpace\MyPython\models\Phi-3-mini-4k-instruct.Q4_K_M.gguf",
    temperature=0.1,
    n_ctx=2048,
    n_batch=512,  # 批处理大小
    verbose=False,  # 不显示详细日志
    n_gpu_layers=0,  # CPU运行（如果有GPU可以设置层数）
)

# 强化版提示词（带示例格式）
prompt = """请严格按照以下模板用中文解释RAG技术：

【定义】
检索增强生成(Retrieval-Augmented Generation)是...

【工作流程】
1. 检索阶段：...
2. 生成阶段：...

【应用场景】
- 场景1：...
- 场景2：...

请直接输出完整内容，不要包含任何问答标记如[/INST]"""

response = llm.create_chat_completion(
    messages=[{
        "role": "user",
        "content": prompt
    }],
    max_tokens=400,
    stop=["\n\n"],  # 用空行作为停止标记
    temperature=0.1  # 更低随机性
)
response2 = llm2.invoke(prompt)

# 后处理清除无效标记
answer = response['choices'][0]['message']['content']
answer = answer.replace("[/INST]", "").strip()

print("===== 修正后的回答 =====")
print(answer)
print("另一种调用方式产生的回答===")
print(response2)